package com.aiwiown.snackmq.broker.handler;

import com.aiwiown.snackmq.broker.auth.AuthorizationService;
import com.aiwiown.snackmq.broker.filter.SimpleTagFilterEngine;
import com.aiwiown.snackmq.cluster.metadata.MetadataManager;
import com.aiwiown.snackmq.common.auth.Action;
import com.aiwiown.snackmq.common.cluster.BrokerId;
import com.aiwiown.snackmq.common.exception.AuthorizationException;
import com.aiwiown.snackmq.common.handler.MessageHandler;
import com.aiwiown.snackmq.common.handler.RequestContext;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import com.aiwiown.snackmq.storage.exception.PartitionNotFoundException;
import com.aiwiown.snackmq.storage.service.PartitionedStorageService;
import com.aiwiown.snackmq.storage.service.ZeroCopyStorageService;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 负责处理来自消费者的 PULL_REQUEST 请求。
 * 支持零拷贝读取模式，提高大量消息读取时的性能。
 */
@Slf4j
public class ConsumeMessageHandler implements MessageHandler {

    /**
     * 分区存储服务，用于传统读取模式
     */
    private final PartitionedStorageService storageService;
    /**
     * 零拷贝存储服务，用于高性能读取
     */
    private final ZeroCopyStorageService zeroCopyStorageService;
    /**
     * 是否启用零拷贝模式
     */
    private final boolean enableZeroCopy;
    /**
     * 授权服务
     */
    private final AuthorizationService authorizationService;

    /**
     * 注入元数据管理器和 Broker 自身 ID
     */
    private final MetadataManager metadataManager;
    /**
     * Broker 自身 ID
     */
    private final BrokerId selfId;
    /**
     * 每次读取的最大消息数量
     */
    private static final int MAX_MESSAGES_PER_PULL = 32;
    /**
     * 实例化一个过滤器引擎
     */
    private final SimpleTagFilterEngine filterEngine = new SimpleTagFilterEngine();

    /**
     * 构造函数，初始化消息处理逻辑。
     *
     * @param storageService         分区存储服务
     * @param zeroCopyStorageService 零拷贝存储服务
     * @param enableZeroCopy         是否启用零拷贝模式
     * @param authorizationService   授权服务
     */
    public ConsumeMessageHandler(PartitionedStorageService storageService,
                                 ZeroCopyStorageService zeroCopyStorageService,
                                 boolean enableZeroCopy, AuthorizationService authorizationService,
                                 MetadataManager metadataManager, BrokerId selfId) {
        this.storageService = storageService;
        this.zeroCopyStorageService = zeroCopyStorageService;
        this.authorizationService = authorizationService;
        this.enableZeroCopy = enableZeroCopy && zeroCopyStorageService != null;
        this.metadataManager = metadataManager;
        this.selfId = selfId;

        if (this.enableZeroCopy) {
            log.info("ConsumeMessageHandler initialized with zero-copy mode enabled");
        } else {
            log.info("ConsumeMessageHandler initialized with traditional read mode");
        }
    }

    @Override
    public void handle(RequestContext ctx, Message request) {
        // 提取请求中的主题、分区ID和偏移量
        String topic = request.getTopic();
        // 【新增】提取消费者组
        String consumerGroup = request.getProperty("consumerGroup");
        int partitionId = Integer.parseInt(request.getProperty("partition", "0"));
        long offset = Long.parseLong(request.getProperty("offset", "0"));
        int maxMessages = Integer.parseInt(request.getProperty("maxMessages", String.valueOf(MAX_MESSAGES_PER_PULL)));
        // 【新增】提取过滤表达式
        String filterExpression = request.getProperty("filterExpression");
        try {
            // --- 【核心优化】授权检查现在包含 Topic 和 ConsumerGroup ---
            authorizationService.authorize(
                    request.getProperty("authToken"),
                    Action.CONSUME,
                    new AuthorizationService.Resource("TOPIC", topic),
                    new AuthorizationService.Resource("CONSUMER_GROUP", consumerGroup)
            );

            TopicPartition topicPartition = new TopicPartition(topic, partitionId);
            // 【修改】重构逻辑以支持过滤
            // 只有在未启用过滤且开启了零拷贝时，才使用零拷贝路径
            boolean useZeroCopy = enableZeroCopy && (filterExpression == null || filterExpression.isEmpty());
            if (useZeroCopy) {
                List<ByteBuffer> messageBuffers = readMessagesWithZeroCopy(topicPartition, offset, maxMessages);
                if (messageBuffers.isEmpty()) {
                    ctx.writeResponse(createNotFoundResponse(request));
                } else {
                    ctx.writeZeroCopyResponse(request.getMessageId(), messageBuffers);
                    log.debug("Successfully sent {} message buffers (zero-copy) from partition {} starting at offset {}",
                            messageBuffers.size(), topicPartition, offset);
                }
            } else {
                // 传统读取路径（有过滤或未开启零拷贝）
                if (enableZeroCopy && filterExpression != null && !filterExpression.isEmpty()) {
                    log.trace("Message filtering is active, falling back to traditional read mode for request on {}", topicPartition);
                }

                List<Message> messages = readMessagesWithTraditionalMode(topicPartition, offset, maxMessages);

                // 【新增】应用过滤逻辑
                final List<Message> filteredMessages;
                if (filterExpression != null && !filterExpression.isEmpty()) {
                    filteredMessages = messages.stream()
                            .filter(msg -> filterEngine.evaluate(msg, filterExpression))
                            .collect(Collectors.toList());
                    log.debug("Applied filter '{}', {} of {} messages matched for partition {}.",
                            filterExpression, filteredMessages.size(), messages.size(), topicPartition);
                } else {
                    filteredMessages = messages;
                }
                // 【修改】使用过滤后的消息列表进行响应
                if (filteredMessages.isEmpty()) {
                    ctx.writeResponse(createNotFoundResponse(request));
                } else {
                    ctx.writeResponse(createSuccessResponse(request, filteredMessages));
                    log.debug("Successfully read and sent {} messages from partition {} starting at offset {}",
                            filteredMessages.size(), topicPartition, offset);
                }
            }
        } catch (AuthorizationException e) {
            log.warn("Authorization failed for pull request: {}", e.getMessage());
            ctx.writeResponse(createErrorResponse(request, "Permission denied: " + e.getMessage()));
        } catch (PartitionNotFoundException e) {
            // --- 【核心修复】: 捕获 PartitionNotFoundException 并进行智能处理 ---
            log.warn("Partition not found locally for pull request: {}. Checking metadata...", request, e);
            TopicPartition tp = new TopicPartition(topic, partitionId);
            // 检查元数据，看这个分区是否 *应该* 在这个 Broker 上
            if (metadataManager.isPartitionAssignedToBroker(tp, selfId.toString())) {
                // 如果是，说明分区正在初始化，这是一个瞬时状态。
                log.info("Partition {} should be on this broker but is not ready. Replying with SERVICE_NOT_READY.", tp);
                sendServiceNotReadyResponse(ctx, request);
            } else {
                // 如果不是，说明客户端的元数据已过时，它请求了一个错误的 Broker。
                // 【修复】: 返回一个包含重定向地址的特定错误响应，而不是一个通用错误。
                String leaderAddress = metadataManager.getPartitionLeaderAddress(tp);
                if (leaderAddress != null) {
                    log.warn("Client requested partition {} which is not on this broker. Redirecting to leader at {}.", tp, leaderAddress);
                    Message redirectResponse = Message.builder()
                            .type(MessageType.RESPONSE)
                            .messageId(request.getMessageId())
                            .status(MessageStatus.FAILED)
                            .errorMessage("Partition not found on this broker. Please redirect to the leader.")
                            .property("redirectAddress", leaderAddress) // 关键的重定向信息
                            .build();
                    ctx.writeResponse(redirectResponse);
                } else {
                    log.error("Partition {} is not on this broker, and no leader could be found in metadata. Replying with failure.", tp);
                    ctx.writeResponse(createErrorResponse(request, "Partition not found and no leader available."));
                }
            }
        } catch (Exception e) {
            log.error("处理 PULL_REQUEST 时出错, topic: {}, partition: {}, offset: {}",
                    topic, partitionId, offset, e);
            ctx.writeResponse(createErrorResponse(request, e));
        }
    }

    /**
     * 【已重构】使用零拷贝模式读取消息。
     * 此方法现在直接返回从存储层获取的 ByteBuffer 切片列表，完成了零拷贝链条在 Handler 层的传递。
     *
     * @param topicPartition 主题分区
     * @param offset         起始偏移量
     * @param maxMessages    最大消息数量
     * @return 消息的 ByteBuffer 列表
     * @throws IOException 如果读取失败
     */
    List<ByteBuffer> readMessagesWithZeroCopy(TopicPartition topicPartition, long offset, int maxMessages)
            throws IOException {
        return zeroCopyStorageService.readZeroCopy(topicPartition, offset, maxMessages);
    }

    /**
     * 【已重构】使用传统模式读取消息。
     * 注意：为了保持一致性，现在也传递 maxMessages 参数。
     *
     * @param topicPartition 主题分区
     * @param offset         起始偏移量
     * @param maxMessages    最大消息数量
     * @return 消息列表
     * @throws IOException 如果读取失败
     */
    List<Message> readMessagesWithTraditionalMode(TopicPartition topicPartition, long offset, int maxMessages)
            throws IOException {
        return storageService.readBatch(topicPartition, offset, maxMessages);
    }

    /**
     * 辅助方法，用于创建响应，提高代码可读性
     *
     * @param request 响应
     * @return 响应
     */
    private Message createNotFoundResponse(Message request) {
        return Message.builder()
                .messageId(request.getMessageId())
                .type(MessageType.RESPONSE)
                .status(MessageStatus.NOT_FOUND)
                .build();
    }

    /**
     * 辅助方法，用于创建响应，提高代码可读性
     *
     * @param request  响应
     * @param messages 消息列表
     * @return 响应
     */
    private Message createSuccessResponse(Message request, List<Message> messages) {
        return Message.builder()
                .messageId(request.getMessageId())
                .type(MessageType.RESPONSE)
                .status(MessageStatus.SUCCESS)
                .messages(messages)
                .build();
    }

    /**
     * 辅助方法，用于创建响应，提高代码可读性
     *
     * @param request      响应
     * @param errorMessage 错误信息
     * @return 响应
     */
    private Message createErrorResponse(Message request, String errorMessage) {
        return Message.builder()
                .messageId(request.getMessageId())
                .type(MessageType.RESPONSE)
                .status(MessageStatus.FAILED)
                .errorMessage(errorMessage)
                .build();
    }

    /**
     * 辅助方法，用于创建响应，提高代码可读性
     *
     * @param request 响应
     * @param e       错误信息
     * @return 响应
     */
    private Message createErrorResponse(Message request, Exception e) {
        return Message.builder()
                .messageId(request.getMessageId())
                .type(MessageType.RESPONSE)
                .status(MessageStatus.FAILED)
                .errorMessage("Failed to read messages from storage: " + e.getMessage())
                .build();
    }

    /**
     * 【已重构】发送 SERVICE_NOT_READY 响应。
     *
     * @param ctx     请求上下文
     * @param request 请求
     */
    private void sendServiceNotReadyResponse(RequestContext ctx, Message request) {
        Message response = Message.builder()
                .type(MessageType.RESPONSE)
                .messageId(request.getMessageId())
                .status(MessageStatus.SERVICE_NOT_READY)
                .errorMessage("Partition is being initialized, please retry shortly.")
                .build();
        ctx.writeResponse(response);
    }

    /**
     * 获取消息类型
     *
     * @return 消息类型
     */
    @Override
    public MessageType getMessageType() {
        return MessageType.PULL_REQUEST;
    }
}